import torch
from torch import nn
from d2l import torch as d2l

def nin_block(in_channels,out_channels,kernel_size,stride,padding):
    net = nn.Sequential(
        nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding),nn.ReLU(),
        nn.Conv2d(out_channels,out_channels,kernel_size=1),nn.ReLU(),
        nn.Conv2d(out_channels,out_channels,kernel_size=1),nn.ReLU(),
    )
    return net

def nin():
    net = nn.Sequential(
        nin_block(1,96,kernel_size=11,stride=4,padding=0),
        nn.MaxPool2d(3,stride=2),
        nin_block(96,256,kernel_size=5,stride=1,padding=2),
        nn.MaxPool2d(3,stride=2),
        nin_block(256,384,kernel_size=3,stride=1,padding=1),
        nn.MaxPool2d(3,stride=2),
        nn.Dropout(),
        nin_block(384,10,kernel_size=3,stride=1,padding=1),
        nn.AdaptiveAvgPool2d((1,1)),
        nn.Flatten()
    )
    return net

